tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md (1)

152-152: ⚡ Quick win

Clarify symmetric-memory provider terminology.

The description states "requires CUDA 13 Cutlass DSL runtime (PR #14354) and NVSHMEM provider". Based on the design doc (phase 0, line 938-939), the implementation uses torch.distributed._symmetric_memory (cuMem-based), not an external NVSHMEM package. Consider revising to "symmetric memory provider" or "torch.distributed._symmetric_memory" to avoid implying a dependency on nvshmem4py-cu13, which was discussed in the design doc (lines 533-534) but not chosen for v1.

📝 Proposed clarification

-| `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR `#14354`) and NVSHMEM provider (hard gate — see `mega_moe/MEGAMOE_CUTEDSL_DESIGN.md`); v1 alpha=1 product gate enforced in `post_load_weights` | `FUSED_COMM` |
+| `mega_moe/mega_moe_cute_dsl.py` | `MegaMoECuteDsl` | SM100/SM103 | NVFP4 via ported CuteDSL `Sm100MegaMoEKernel` fused dispatch+FC1+act+FC2+combine kernel; requires CUDA 13 Cutlass DSL runtime (PR `#14354`) and symmetric memory provider (`torch.distributed._symmetric_memory`); v1 alpha=1 product gate enforced in `post_load_weights` | `FUSED_COMM` |

Based on learnings from MEGAMOE_CUTEDSL_DESIGN.md phase 0: the implementation uses PyTorch's torch.distributed._symmetric_memory (cuMem-based NVSHMEM-equivalent), not an external NVSHMEM package.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md` at line 152,
Update the README entry describing mega_moe/mega_moe_cute_dsl.py /
MegaMoECuteDsl to avoid implying an external NVSHMEM package: change "NVSHMEM
provider" to "symmetric memory provider (torch.distributed._symmetric_memory)"
or similar wording that explicitly names PyTorch's cuMem-based provider; ensure
the note about requiring CUDA 13 Cutlass DSL runtime and the product gate
references remain unchanged so readers know runtime and gating requirements.

tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py (1)

706-706: 💤 Low value

Consider ClassVar annotation for shared mutable class attribute.

The kernel_cache is intentionally shared across all instances for compile caching. Adding ClassVar makes this intent explicit and silences the RUF012 warning:

+from typing import ClassVar
+
 class Sm100MegaMoENvfp4Runner(TunableRunner):
     ...
     # Module-scope compile cache shared by every runner instance.
-    kernel_cache: dict = {}
+    kernel_cache: ClassVar[dict] = {}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py` at line 706,
The class-level mutable cache kernel_cache is intended to be shared across
instances but lacks an explicit ClassVar annotation; update the declaration of
kernel_cache in the class to be typed as typing.ClassVar[dict] (and add from
typing import ClassVar if not already imported) so the intent is explicit and
the RUF012 warning is silenced while preserving the shared compile cache
behavior.

tests/unittest/_torch/modules/moe/test_moe_backend.py (2)

568-577: ⚡ Quick win

Add a companion negative test for the v1 alpha gate.

Setting TRTLLM_MEGAMOE_CUTEDSL_BYPASS_V1_ALPHA_GATE=1 for this path is useful for end-to-end execution, but it also removes direct coverage that the default gate rejects non-1 alpha checkpoints. Please add a small paired test with the env var unset that asserts rejection.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 568 -
577, Add a companion negative test that mirrors the MEGAMOE_CUTEDSL branch but
ensures the v1 alpha gate is enforced: when backend_type ==
MoeBackendType.MEGAMOE_CUTEDSL, explicitly ensure TR
TLLM_MEGAMOE_CUTEDSL_BYPASS_V1_ALPHA_GATE is unset (use monkeypatch.delenv or
monkeypatch.setenv with None), then exercise the same load -> post-load ->
run_moe path used in the positive case and assert that the load/post-load step
fails (raises/rejects) due to non-1 alpha values; place this next to the
existing positive branch in
tests/unittest/_torch/modules/moe/test_moe_backend.py so it uses the same setup
and failure assertion to prove the gate blocks non-1 alpha checkpoints.

---

307-315: QA integration list updates are unnecessary for this change-set.

This file only expands unit-test backend coverage; no tests/integration/defs/ entries were added or materially changed, so tests/integration/test_lists/qa/* updates are not required in this PR.

As per coding guidelines: "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 307 -
315, Add a brief explicit note that QA integration list updates are unnecessary
because this change only expands unit-test backend coverage and does not modify
tests/integration/defs; place the comment immediately above
BACKEND_TYPES_TO_TEST (referencing BACKEND_TYPES_TO_TEST, MoeBackendType.*
entries) so reviewers and release/QA scripts know no
tests/integration/test_lists/qa/* edits are required.

tests/microbenchmarks/bench_moe/utils.py (1)

66-79: 💤 Low value

_ensure_dist_for_megamoe only handles MEGAMOE_DEEPGEMM, not MEGAMOE_CUTEDSL.

This function checks only for MEGAMOE_DEEPGEMM but test_moe_module.py:_ensure_dist_for_megamoe (lines 139-143) checks for both MegaMoE variants. If bench_moe is ever used with MEGAMOE_CUTEDSL, the distributed process group won't be initialized.

Since MEGAMOE_CUTEDSL isn't in the benchmark's MoeBackendType enum yet (preview phase), this is acceptable for now but should be updated when CuteDSL is added to the benchmark registry.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/microbenchmarks/bench_moe/utils.py` around lines 66 - 79, The
early-return in _ensure_dist_for_megamoe only recognizes
MoeBackendType.MEGAMOE_DEEPGEMM, so the NCCL ProcessGroup will not be
initialized for the other MegaMoE variant; update _ensure_dist_for_megamoe to
also accept the MEGAMOE_CUTEDSL backend (or generally any future MegaMoE enum
value) by checking for both MoeBackendType.MEGAMOE_DEEPGEMM.value and
MoeBackendType.MEGAMOE_CUTEDSL.value (or by matching a common prefix/enum
category), keeping the rest of the initialization (CUDA check, env vars,
dist.init_process_group) unchanged; reference function name
_ensure_dist_for_megamoe and the test file
test_moe_module.py:_ensure_dist_for_megamoe to ensure parity with tests and add
a TODO comment to revisit when CuteDSL is formally added to the benchmark
registry.

tests/microbenchmarks/bench_moe/backend.py (1)

121-124: 💤 Low value

Consider adding MEGAMOE_CUTEDSL to the benchmark registry.

The MoeBackendType enum and get_backend_class dispatch only include MEGAMOE_DEEPGEMM. The CuteDSL variant is missing, which means bench_moe cannot benchmark the new backend.

This is acceptable for the preview phase, but should be addressed before the feature graduates to production.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/microbenchmarks/bench_moe/backend.py` around lines 121 - 124, Add the
missing CuteDSL backend to the registry by extending the MoeBackendType enum
with a MEGAMOE_CUTEDSL member and updating the get_backend_class dispatch to
handle it: import the CuteDSL backend implementation (e.g., the MegaMoE CuteDSL
class) alongside MegaMoEDeepGemm and return that class when backend_type ==
MoeBackendType.MEGAMOE_CUTEDSL; ensure you reference the new enum value
(MoeBackendType.MEGAMOE_CUTEDSL) and the CuteDSL class name (e.g.,
MegaMoECuteDSL) in the dispatch so bench_moe can instantiate the CuteDSL
backend.

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py (2)

1926-1927: ⚡ Quick win

Prefix unused variable with underscore.

The n_idx variable from the CLC response unpacking is never used. Prefix it with underscore to indicate it's intentionally ignored.

Suggested fix

-        m_idx, n_idx, l_idx, is_valid = cute.arch.clc_response(
+        m_idx, _n_idx, l_idx, is_valid = cute.arch.clc_response(

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py`
around lines 1926 - 1927, The unpacking of the CLC response in the call to
cute.arch.clc_response currently assigns an unused variable n_idx; change that
identifier to a prefixed name (e.g., _n_idx) so it signals intentional ignore.
Locate the line where m_idx, n_idx, l_idx, is_valid =
cute.arch.clc_response(...) (inside moe_persistent_scheduler.py) and replace
n_idx with _n_idx in the tuple assignment; leave m_idx, l_idx and is_valid
unchanged.

---

7-7: 💤 Low value

Consider using built-in types instead of typing module.

Per coding guidelines, prefer built-in list, tuple over typing.List, typing.Tuple for Python 3.10+. Also prefer X | None over Optional[X].

This is a low-priority cleanup given the file size.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_persistent_scheduler.py`
at line 7, Replace typing.List, typing.Tuple, and typing.Optional usages with
built-in generic types and PEP 604 unions: change annotations like List[X] ->
list[X], Tuple[A, B] -> tuple[A, B], and Optional[T] -> T | None; remove List,
Optional, Tuple from the import line in moe_persistent_scheduler.py and
keep/import only Literal if still needed; update all function/method signatures
and variable annotations that reference List, Tuple, or Optional accordingly
(search for usages of List, Tuple, Optional in the file and replace them with
list, tuple, and | None forms).

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/blocked_scale.py (1)

33-33: 💤 Low value

Prefer built-in list over typing.List.

Per coding guidelines, prefer built-in types like list over legacy typing.List for Python 3.10+.

Suggested fix

-from typing import List
+from collections.abc import Sequence

Then update function signatures:

-def cat_byte_reinterpretable_tensors(tensors: List[torch.Tensor], dim: int = 0) -> torch.Tensor:
+def cat_byte_reinterpretable_tensors(tensors: list[torch.Tensor], dim: int = 0) -> torch.Tensor:

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/blocked_scale.py` at line
33, Replace the legacy typing.List import and usages with the built-in list
type: remove "from typing import List" and update all annotations in this module
(e.g., function signatures, return types, and variable annotations that
reference List) to use the native list[...] form (for example change "List[T]"
to "list[T]"); ensure functions/methods in this file that reference List (search
for symbols like any function or class definitions using List) are updated
accordingly and the typing import is deleted.

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py (1)

5-13: 💤 Low value

Consider UPPER_SNAKE_CASE for module-level constants.

Per coding guidelines, Python constants should use UPPER_SNAKE_CASE (e.g., NVFP4_BLOCK_SIZE, SF_PADDING_BLOCK). The current PascalCase naming is internally consistent but deviates from the project convention.

Given this is a preview PR, this can be deferred if the naming aligns with an upstream source being ported.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py`
around lines 5 - 13, Module-level constants use PascalCase; rename them to
UPPER_SNAKE_CASE to match project conventions: change Nvfp4BlockSize ->
NVFP4_BLOCK_SIZE, SfPaddingBlock -> SF_PADDING_BLOCK, TmaLeadingDimByteAlign ->
TMA_LEADING_DIM_BYTE_ALIGN, Nvfp4E2M1Max -> NVFP4_E2M1_MAX, Fp8E4M3FNMax ->
FP8_E4M3_FN_MAX, SupportedMmaTileM -> SUPPORTED_MMA_TILE_M, SupportedMmaTileN ->
SUPPORTED_MMA_TILE_N, and update all internal references/usages of these symbols
accordingly; if external consumers may rely on the old names, add short-lived
aliases mapping the old names to the new constants in the same module to
preserve backward compatibility (e.g., Nvfp4BlockSize = NVFP4_BLOCK_SIZE) and
add a TODO to remove aliases later.

tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py (1)

100-100: 💤 Low value

Modernize type hints to use built-in types.

The coding guidelines prefer built-in types (list, dict, tuple) over legacy typing equivalents, and X | None over Optional[X]. This file imports and uses List, Dict, Tuple, Optional throughout.

♻️ Suggested import and usage update

-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any

Then update usages throughout, e.g.:

• List[ir.Value] → list[ir.Value]
• Dict[str, int] → dict[str, int]
• Tuple[int, ...] → tuple[int, ...]
• Optional[int] → int | None
• Type[cutlass.Numeric] → type[cutlass.Numeric]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py` at
line 100, Replace legacy typing generics with built-in generics and PEP 604
unions: remove List, Dict, Tuple, Optional, Type from the typing import and use
list, dict, tuple, X | None, and type[...] in their places; e.g., change imports
on the line with "from typing import Any, Dict, List, Optional, Tuple, Type" to
keep only Any (if needed) and update all occurrences like
List[ir.Value]→list[ir.Value], Dict[str,int]→dict[str,int],
Tuple[int,...]→tuple[int,...], Optional[int]→int | None, and
Type[cutlass.Numeric]→type[cutlass.Numeric] throughout megamoe_kernel.py (search
for those symbols: List, Dict, Tuple, Optional, Type and replace accordingly).

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes